#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
使用阿里云 Qwen 模型的数据可视化引擎示例
"""

import asyncio
import pandas as pd
import json
import os
from datetime import datetime
from data_engine.core.engine import DataEngine
from data_engine.utils.llm_adapter import LLMClient
from data_engine.utils.logger import get_logger, setup_logging

# 设置日志级别为DEBUG
setup_logging(level="DEBUG")

# 获取测试日志器
test_logger = get_logger("qwen_example")


async def main():
    # 初始化 LLM 客户端
    llm_client = LLMClient(
        model="openai/qwen-turbo",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
        api_key="sk-kkkkkkkkkkkkkkkkkkk"  
    )
    
    # 初始化数据引擎并启用调试模式
    engine = DataEngine(llm_client, debug_mode=True)
    
    # 加载销售数据（使用UTF-8编码）
    sales_data = pd.read_csv("test_data/sales_data.csv", encoding='utf-8')
    field_descriptions = {
        "month": "月份",
        "sales_amount": "销售金额（元）",
        "department": "部门",
        "region": "地区"
    }
    
    # 加载数据集
    engine.load_dataset(
        df=sales_data,
        dataset_desc="公司销售数据",
        table_desc="记录各部门各地区每月销售金额",
        field_desc=field_descriptions
       
    )
    
    # 示例查询
    question = "计算每个部门的平均销售额，显示平均销售额大于6万的部门"
    test_logger.info(f"分析问题: {question}")
    
    try:
        result = await engine.process_query(question)
        
        # 准备保存的JSON数据
        result_data = {
            "timestamp": datetime.now().isoformat(),
            "question": question,
            "success": result.success,
            "chart_type": result.chart_type,
            "execution_time": result.execution_time,
            "data_rows": len(result.data) if result.data is not None else 0,
            "chart_config": result.chart_config,
            "error": result.error if not result.success else None,
            "analysis_result": result.analysis_result,
            "data_sample": result.data[:5] if result.data is not None and len(result.data) > 0 else None
        }
        
        # 确保test_output目录存在
        os.makedirs("test_output", exist_ok=True)
        
        # 生成文件名（使用时间戳）
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"test_output/result_{timestamp}.json"
        
        # 保存JSON结果
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(result_data, f, ensure_ascii=False, indent=2)
        
        test_logger.info(f"📁 结果已保存到: {filename}")
        
        if result.success:
            test_logger.info("✅ 分析成功")
            test_logger.info(f"图表类型: {result.chart_type}")
            test_logger.info(f"处理时间: {result.execution_time:.2f}秒")
            test_logger.info(f"数据行数: {len(result.data)}")
            
            # 显示配置概要
            if result.chart_config:
                test_logger.info(f"图表配置包含: {list(result.chart_config.keys())}")
                # 显示部分配置信息
                if 'title' in result.chart_config:
                    test_logger.info(f"图表标题: {result.chart_config.get('title', {}).get('text', 'N/A')}")
                if 'xAxis' in result.chart_config:
                    test_logger.info(f"X轴配置: 已设置")
                if 'yAxis' in result.chart_config:
                    test_logger.info(f"Y轴配置: 已设置")
                if 'series' in result.chart_config:
                    series_count = len(result.chart_config.get('series', []))
                    test_logger.info(f"数据系列: {series_count}个")
                
                # 输出完整的ECharts配置JSON
                test_logger.info("\n📈 ECharts图表配置JSON:")
                test_logger.info(json.dumps(result.chart_config, ensure_ascii=False, indent=2))
        else:
            test_logger.error(f"❌ 分析失败: {result.error}")
            
    except Exception as e:
        test_logger.error(f"❌ 发生异常: {str(e)}")
        # 即使发生异常也保存错误信息
        error_data = {
            "timestamp": datetime.now().isoformat(),
            "question": question,
            "success": False,
            "error": str(e),
            "exception_type": type(e).__name__
        }
        
        os.makedirs("test_output", exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"test_output/error_{timestamp}.json"
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(error_data, f, ensure_ascii=False, indent=2)
        
        test_logger.error(f"📁 错误信息已保存到: {filename}")


if __name__ == "__main__":
    asyncio.run(main())